import torch.nn as nn
import torch.nn.functional as F

# LeNet
class CNN(nn.Module):
    def __init__(self, num_classes):
        nn.Module.__init__(self)
        self.conv1 = nn.Sequential(
            # Input shape (3, 64, 64)
            nn.Conv2d(
                in_channels=3,
                out_channels=6,
                kernel_size=5,
                stride=1,
                padding=2
            ),
            # Output shape (6, 60, 60)
            nn.ReLU(),
            # Output shape (6, 30, 30)
            nn.MaxPool2d(kernel_size=2)
        )
        self.conv2 = nn.Sequential(
            # Input shape (6, 30, 30)
            nn.Conv2d(
                in_channels=6,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2
            ),
            # Output shape (16, 26, 26)
            nn.ReLU(),
            # Output shape (16, 13, 13)
            nn.MaxPool2d(kernel_size=2)
        )
        self.fc = nn.Sequential(
            # FC output 5 classes
            nn.Linear(in_features=16 * 16 * 16,
                      out_features=300),
            nn.ReLU(),
            nn.Linear(in_features=300,
                      out_features=84),
            nn.ReLU(),
            nn.Linear(in_features=84,
                      out_features=num_classes)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size()[0], -1)
        x = self.fc(x)
        return x
